#!/usr/bin/env python3
import rospy
import numpy as np
import open3d as o3d
import os
from datetime import datetime
from sensor_msgs.msg import PointCloud2
import sensor_msgs.point_cloud2 as pc2

class PointCloudSaver:
    def __init__(self):
        rospy.init_node('pointcloud_saver', anonymous=True)
        
        # 参数设置
        self.save_path = rospy.get_param('~save_path', os.path.expanduser('/root/ros_ws/src/OrbbecSDK_ROS1/scripts/data/pointclouds'))
        self.max_saves = rospy.get_param('~max_saves', 5)  # 最多保存的点云数
        self.topic_name = rospy.get_param('~topic_name', '/camera/depth_registered/points')  # 支持动态话题
        self.voxel_size = rospy.get_param('~voxel_size', 0.01)  # 降采样体视素大小（米）
        self.z_min = rospy.get_param('~z_min', 0.1)  # 最小深度（米）
        self.z_max = rospy.get_param('~z_max', 6.0)  # 最大深度（米）
        self.frame_count = 0
        
        # 创建保存目录
        os.makedirs(self.save_path, exist_ok=True)
        rospy.loginfo(f"点云将保存至: {self.save_path}")
        rospy.loginfo(f"订阅点云话题: {self.topic_name}")
        rospy.loginfo(f"体视素降采样大小: {self.voxel_size} 米")
        rospy.loginfo(f"深度范围: [{self.z_min}, {self.z_max}] 米")
        
        # 订阅点云话题
        self.cloud_sub = rospy.Subscriber(
            self.topic_name,
            PointCloud2,
            self.cloud_callback,
            queue_size=1
        )
        
        rospy.loginfo("等待点云数据...")

    def cloud_to_o3d(self, msg):
        """将 ROS PointCloud2 转换为 Open3D 点云"""
        # 打印字段信息
        rospy.loginfo(f"点云字段: {str(msg.fields)}")
        
        # 提取 XYZ 数据
        points = np.array(list(pc2.read_points(msg, skip_nans=True, field_names=("x", "y", "z"))))
        if len(points) == 0:
            rospy.logwarn("收到空点云，跳过")
            raise ValueError("点云数据为空")
        
        # 统计原始点数
        total_points = len(points)
        
        # 过滤无效点
        valid_mask = (points[:, 2] > self.z_min) & (points[:, 2] < self.z_max) & \
                     (~np.isnan(points[:, 0])) & (~np.isnan(points[:, 1])) & (~np.isnan(points[:, 2]))
        points = points[valid_mask]
        
        # 统计无效点原因
        valid_points = len(points)
        invalid_z_min = np.sum(points[:, 2] <= self.z_min) if len(points) > 0 else 0
        invalid_z_max = np.sum(points[:, 2] >= self.z_max) if len(points) > 0 else 0
        invalid_nan = total_points - valid_points
        rospy.loginfo(f"原始点数: {total_points}, 有效点数: {valid_points}, 有效比例: {valid_points/total_points:.2%}")
        rospy.loginfo(f"无效点统计: z <= {self.z_min}: {invalid_z_min}, z >= {self.z_max}: {invalid_z_max}, nan: {invalid_nan}")
        
        if valid_points == 0:
            rospy.logwarn("所有点无效，跳过")
            raise ValueError("所有点无效")
        
        # 打印点云坐标范围
        rospy.loginfo(f"点云坐标范围: x=[{points[:, 0].min():.2f}, {points[:, 0].max():.2f}], "
                      f"y=[{points[:, 1].min():.2f}, {points[:, 1].max():.2f}], "
                      f"z=[{points[:, 2].min():.2f}, {points[:, 2].max():.2f}]")
        
        # 创建 Open3D 点云
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points)
        
        # 处理颜色数据（rgb 为 float32）
        if 'rgb' in str(msg.fields):
            try:
                colors = np.array(list(pc2.read_points(msg, skip_nans=True, field_names=("rgb",))))
                colors = colors.flatten()[valid_mask]
                if len(colors) == len(points):
                    colors = colors.astype(np.float32)
                    rospy.loginfo(f"RGB 数据范围: [{colors.min():.2f}, {colors.max():.2f}]")
                    rgb = np.zeros((len(colors), 3))
                    # 尝试多种解码方式
                    if colors.max() <= 1.0:  # 0-1 范围
                        rgb[:, 0] = np.clip(colors, 0, 1)
                        rgb[:, 1] = np.clip(colors, 0, 1)
                        rgb[:, 2] = np.clip(colors, 0, 1)
                    elif colors.max() <= 255.0:  # 0-255 范围
                        rgb[:, 0] = np.clip(colors, 0, 255) / 255.0
                        rgb[:, 1] = np.clip(colors, 0, 255) / 255.0
                        rgb[:, 2] = np.clip(colors, 0, 255) / 255.0
                    else:  # 打包格式 r*256*256 + g*256 + b
                        rgb[:, 0] = np.clip(colors / (256.0 * 256.0), 0, 255) / 255.0
                        rgb[:, 1] = np.clip((colors / 256.0) % 256.0, 0, 255) / 255.0
                        rgb[:, 2] = np.clip(colors % 256.0, 0, 255) / 255.0
                    pcd.colors = o3d.utility.Vector3dVector(rgb)
                    rospy.loginfo(f"处理颜色数据，点数: {len(colors)}")
                else:
                    rospy.logwarn("颜色数据和点云数量不匹配，跳过颜色")
            except Exception as e:
                rospy.logwarn(f"颜色处理失败: {str(e)}，跳过颜色")
        
        # 体视素降采样
        if self.voxel_size > 0:
            pcd = pcd.voxel_down_sample(voxel_size=self.voxel_size)
            rospy.loginfo(f"降采样后点数: {len(pcd.points)}")
        
        return pcd

    def cloud_callback(self, msg):
        """点云回调函数 - 转换并保存点云"""
        try:
            # 转换点云为 Open3D 格式
            pcd = self.cloud_to_o3d(msg)
            self.frame_count += 1
            
            # 生成带时间戳的文件名
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = os.path.join(self.save_path, f"orbbec_{timestamp}.ply")
            
            # 保存点云为 PLY 格式
            o3d.io.write_point_cloud(filename, pcd)
            rospy.loginfo(f"成功保存点云至: {filename}")
            
            # 删除旧文件以控制存储空间
            self.cleanup_old_files()
            
            # 达到最大保存数量后自动关闭
            if self.frame_count >= self.max_saves:
                rospy.loginfo(f"已达到最大保存数量({self.max_saves})，关闭节点")
                rospy.signal_shutdown("Completed")
                
        except Exception as e:
            rospy.logerr(f"点云处理失败: {str(e)}")
        
    def cleanup_old_files(self):
        """删除旧文件以控制存储空间"""
        if self.frame_count <= self.max_saves * 2:
            return
            
        # 获取目录中的所有点云文件
        pc_files = [f for f in os.listdir(self.save_path) if f.endswith('.ply')]
        if len(pc_files) <= self.max_saves:
            return
            
        # 按修改时间排序
        pc_files.sort(key=lambda x: os.path.getmtime(os.path.join(self.save_path, x)))
        
        # 删除最旧的文件，只保留最近的 max_saves 个文件
        files_to_delete = pc_files[:-self.max_saves]
        for f in files_to_delete:
            file_path = os.path.join(self.save_path, f)
            try:
                os.remove(file_path)
                rospy.loginfo(f"已删除旧文件: {file_path}")
            except Exception as e:
                rospy.logerr(f"删除文件失败 {file_path}: {str(e)}")

    def run(self):
        """主循环"""
        try:
            rospy.spin()
        except rospy.ROSInterruptException:
            rospy.loginfo("节点被中断")

if __name__ == "__main__":
    try:
        saver = PointCloudSaver()
        saver.run()
    except Exception as e:
        rospy.logerr(f"初始化失败: {str(e)}")
